import cv2
import random
import numpy as np

def rand(a=0,b=1):
    return np.random.rand()*(b-a) + a

def letterbox_v1(batch_images, batch_annotations, input_shape=(640,640)):
    w, h = input_shape
    for idx in range(len(batch_images)):
        raw_image = batch_images[idx].copy()
        gen_image = np.zeros(shape=(h,w,3))
        raw_annotations = batch_annotations[idx].copy()
        rh, rw, rc = raw_image.shape
        r = rw / rh
        if r > 1:
            nw = w
            nh = int(nw / r + 0.5)
            dw = 0
            dh = (h - nh) // 2
        else:
            nh = h
            nw = int(nh * r + 0.5)
            dw = (w - nw) // 2
            dh = 0
        raw_image = cv2.resize(raw_image,(nw,nh),interpolation=cv2.INTER_LINEAR)
        gen_image[dh:dh+nh,dw:dw+nw,:] = raw_image
        raw_annotations[:,1] = dw / w + raw_annotations[:,1] * nw / w
        raw_annotations[:,2] = dh / h + raw_annotations[:,2] * nh / h
        raw_annotations[:,3] = raw_annotations[:,3] * nw / w
        raw_annotations[:,4] = raw_annotations[:,4] * nh / h

        batch_images[idx] = gen_image
        batch_annotations[idx] = raw_annotations
        
    return batch_images, batch_annotations

def letterbox_v2(batch_images, batch_annotations, input_shape=(640,640), color=128, auto=True, strides=32):
    # ! input_shape: (w,h)
    for idx in range(len(batch_images)):
        raw_image = batch_images[idx].copy() # ! 原始图片
        rh, rw, rc = raw_image.shape         # ! 原始图片尺寸
        r = min(input_shape[1]/rw, input_shape[0]/rh) # ! 需要尺寸与原图尺寸宽高分别的比例
        nh, nw = int(round(rh * r)), int(round(rw * r))
        dh, dw = (nh - rh), (nw - rw)
        if auto:
            dh, dw = dh % strides, dw % strides # ! 最小矩形
        if (nw,nh) != (rw,rh): # ! scale有变化
            raw_image = cv2.resize(raw_image,(nw,nh),interpolation=cv2.INTER_LINEAR)
    
        gen_image = np.ones(shape=(nh+dh, nw+dw, 3)) * color  # ! 拟生成图片
        dh, dw = dh // 2, dw // 2
        gen_image[dh:dh+nh,dw:dw+nw,:] = raw_image

        if batch_annotations is not None:
            raw_annotations = batch_annotations[idx].copy() # ! 原始标注
            raw_annotations[:,1] = dw / (nw+2*dw) + raw_annotations[:,1] * nw / (nw+2*dw)
            raw_annotations[:,2] = dh / (nh+2*dh) + raw_annotations[:,2] * nh / (nh+2*dh)
            raw_annotations[:,3] = raw_annotations[:,3] * nw / (nw+2*dw)
            raw_annotations[:,4] = raw_annotations[:,4] * nh / (nh+2*dh)

        batch_images[idx] = gen_image
        batch_annotations[idx] = raw_annotations
        
    return batch_images, batch_annotations

def mosaic(batch_images,batch_annotations,jitter=0.3):
    # !建议t取值在[0.04,0.11]之间 [0.2+0.04,0.2+0.11]
    t = rand(0.04,0.31)

    #! 图片数据增强
    for idx in range(len(batch_images)):
        raw_image = batch_images[idx].copy()
        raw_annotations = batch_annotations[idx].copy()
        #! image默认使用cv2读取, 通道格式默认为BGR, shape形状为(h,w,c)
        rh, rw, rc = raw_image.shape
        #! 需要在annotations完成向model预测结果转换之前完成图像增强处理

        #! 翻转
        flip = rand() < t
        if flip and len(raw_annotations) > 0:
            raw_image, raw_annotations = flip_together(raw_image,raw_annotations)

        #! 旋转
        rotate = rand() < t
        if rotate and len(raw_annotations) > 0:
            raw_image, raw_annotations = rotate_together(raw_image,raw_annotations)
            
        #! 扭曲|相对位置不变|再缩放贴到原图大小的画布
        distort = rand() < t
        if distort and len(raw_annotations) > 0:
            raw_image, raw_annotations = distort_together(raw_image,raw_annotations,jitter)

        #! 图片通道变换|标注不需要调整
        channel = rand() < t
        if channel and len(raw_annotations)>0: 
            raw_image = channel_transform(raw_image)

        #! 图片色域变换|标注不需要调整
        gamut = rand() < t
        if gamut and len(raw_annotations) > 0:
            raw_image = gamut_transfrom(raw_image)
        
        batch_images[idx] = raw_image
        batch_annotations[idx] = raw_annotations

    return batch_images, batch_annotations

def flip_together(image,annotations):
    raw_image = image.copy()
    raw_annotations = annotations.copy()
    flip_horizontal = rand() >= 0.5
    if flip_horizontal: #! 水平翻转, y轴不变
        raw_image = cv2.flip(raw_image,1)
        raw_annotations[:,1] = 1 - raw_annotations[:,1] # ! 0为类别
    else:               #! 垂直翻转, x轴不变
        raw_image = cv2.flip(raw_image,0)
        raw_annotations[:,2] = 1 - raw_annotations[:,2]
    return raw_image, raw_annotations

def rotate_together(image,annotations):
    raw_image = image.copy()
    raw_annotations = annotations.copy()

    #! 2代表cv2.ROTATE_90_COUNTERCLOCKWISE|逆时针90度 |顺时针270度
    #! 1代表cv2.ROTATE_180                |逆时针180度|顺时针180度
    #! 0代表cv2.ROTATE_90_CLOCKWISE       |逆时针270度|顺时针90度    

    # ! 图像旋转
    direction = random.choice([2,1,0])
    if direction==2:
        angle = 270
        raw_image = cv2.rotate(raw_image,cv2.ROTATE_90_CLOCKWISE)
    elif direction==1:
        angle = 180
        raw_image = cv2.rotate(raw_image,cv2.ROTATE_180)
    else: 
        angle = 90
        raw_image = cv2.rotate(raw_image,cv2.ROTATE_90_COUNTERCLOCKWISE) # 逆时针90度

    # ! 旋转后坐标处理
    radian = np.math.pi * angle / 180
    xyc = raw_annotations[:,1:3] - 0.5 # ! 以(0.5,0.5)重新建立坐标系, 中心点的位置
    M = np.array([ # ! 图像坐标系, 逆时针旋转
        [ np.math.cos(radian), -np.math.sin(radian)],
        [ np.math.sin(radian),  np.math.cos(radian)]
    ])
    new_xyc = xyc @ M + 0.5
    raw_annotations[:,1:3] = new_xyc

    if angle == 90 or angle == 270: # ! 宽高交换 ?? 为什么要做宽高交换
        raw_annotations[:,[3,4]] = raw_annotations[:,[4,3]]

    return raw_image, raw_annotations

def distort_together(image,annotations,jitter,color=128):
    # ! 扭曲后的图片尺寸与原图片尺寸一致, 扭曲后用空白条填充
    raw_image = image.copy()
    raw_annotations = annotations.copy()
    rh, rw, rc = raw_image.shape

    #! 对图片进行扭曲
    r = rw/rh * rand(1-jitter,1+jitter)/rand(1-jitter,1+jitter) #! 得到新的宽高比
    scale = rand(0.4,1)
    if r < 1: #! 扭曲后的图h更大
        sh = int(round(rh * scale))
        sw = int(round(sh * r))
    else:     #! 扭曲后的图w更大
        sw = int(round(rw * scale))
        sh = int(round(sw / r))
    #! 如果新图片大小超原图宽高
    if sh > rh:
        scale = rand(0.4,rh/sh)
        sh = int(round(sh * scale))
        sw = int(round(sw * scale))
    elif sw > rw:
        scale = rand(0.4,rw/sw)
        sw = int(round(sw * scale))
        sh = int(round(sh * scale))
    raw_image = cv2.resize(raw_image,(sw,sh),interpolation=cv2.INTER_LINEAR) # ! 缩放至新尺寸

    # ! 生成背景图
    if color is not None:
        gen_image = np.ones(shape=(rh,rw,3),dtype=np.uint8) * color
    else:
        pure_color = rand() < 0.5
        if pure_color: #! 纯色模式
            gen_image = np.ones(shape=(rh,rw,3),dtype=np.uint8) # ! 生成图片
            gen_image[...,0] = np.random.randint(0,256) #! np.random.randint(a,b)生成的数不包含b
            gen_image[...,1] = np.random.randint(0,256)
            gen_image[...,2] = np.random.randint(0,256)
        else: #! 杂色模式
            gen_image = np.random.randint(0,256,size=(rh,rw,3),dtype=np.uint8)

    #! 进行贴图并处理标注
    dwr = rand(0.33,0.67)
    dhr = rand(0.33,0.67)
    dw = int(round((rw - sw) * dwr))
    dh = int(round((rh - sh) * dhr))
    gen_image[dh:dh+sh,dw:dw+sw] = raw_image
    raw_annotations[:,1] = dw / rw + raw_annotations[:,1] * sw / rw
    raw_annotations[:,2] = dh / rh + raw_annotations[:,2] * sh / rh
    raw_annotations[:,3] = raw_annotations[:,3] * sw / rw
    raw_annotations[:,4] = raw_annotations[:,4] * sh / rh

    return gen_image, raw_annotations

def channel_transform(image):
    raw_image = image.copy()
    #! 0:BGR|012, 1:BRG|021, 2:GRB|120, 3:GBR|102, 4:RGB|210, 5:RBG|201
    mode = random.choice([0,1,2,3,4,5])
    if mode==5:
        raw_image = raw_image[...,[2,0,1]]
    elif mode==4:
        raw_image = raw_image[...,[2,1,0]]
    elif mode==3:
        raw_image = raw_image[...,[1,0,2]]
    elif mode==2:
        raw_image = raw_image[...,[1,2,0]]
    elif mode==1:
        raw_image = raw_image[...,[0,2,1]]
    else:
        pass
    return raw_image

def gamut_transfrom(image):
    raw_image = image.copy()
    raw_image = cv2.cvtColor(raw_image,cv2.COLOR_BGR2HSV)
    #! 0:HSV, 1:HVS, 2:SVH, 3:SHV, 4:VSH, 5:VHS
    #! 0:HSV|012, 1: HVS|021, 其他情况会出现值超出范围
    mode = random.choice([0,1]) # ! 模式
    if mode:
        raw_image = raw_image[...,[0,2,1]]
        choice = random.choice([0,1]) # ! 转回BGR
        if choice:
            raw_image = cv2.cvtColor(raw_image,cv2.COLOR_HSV2BGR)

    return raw_image

def makeup(batch_images,batch_annotations,num_ratio=0.25):
    #! 用融合图替换普通图
    batch = len(batch_images)
    indexs = np.arange(batch)
    sub_num = random.randint(0,int(round(batch * num_ratio))) #! 需要替换的数量, 不超过1/4
    sub_indexs = np.unique(np.random.choice(indexs,size=sub_num)) #! 去重之后 sub_indexs的size可能小于sub_num
    sub_images = []
    sub_annotations = []
    for sisi in range(sub_indexs.size):
        #! 选择4张图片进行融合
        iis = np.unique(np.random.choice(indexs,size=4)) # image indexes
        while iis.size < 4:
            temp = iis[-1]
            if temp+1==batch: temp=0
            iis = np.concatenate([iis,[temp+1]])
        mh,mw,mc = np.max([batch_images[i].shape for i in iis],axis=0) #! 取最大的h,w,c
        cx = int(mw * rand(0.3,0.7)) # crossover x
        cy = int(mh * rand(0.3,0.7)) # crossover y
        gen_image = np.zeros(shape=(mh,mw,mc),dtype=np.uint8)
        gen_annotations = []
        for loc,ii in enumerate(iis): # image index
            #! 0
            raw_image = batch_images[ii].copy()
            raw_annotations = batch_annotations[ii].copy()
            rh, rw, rc = raw_image.shape
            if loc == 0: #! 左上
                nw, nh = cx, cy
                dw, dh =  0,  0
            elif loc == 1: #! 右上
                nw, nh = mw - cx, cy
                dw, dh =      cx,  0
            elif loc == 2: #! 左下
                nw, nh = cx, mh - cy
                dw, dh =  0,      cy
            else: #! 右下
                nw, nh = mw - cx, mh - cy
                dw, dh =      cx,      cy
            raw_image = cv2.resize(raw_image,(nw,nh),interpolation=cv2.INTER_LINEAR)
            gen_image[dh:dh+nh,dw:dw+nw] = raw_image
            raw_annotations[:,1] = dw / mw + raw_annotations[:,1] * nw / mw
            raw_annotations[:,2] = dh / mh + raw_annotations[:,2] * nh / mh
            raw_annotations[:,3] = raw_annotations[:,3] * nw / mw
            raw_annotations[:,4] = raw_annotations[:,4] * nh / mh   
            gen_annotations.append(raw_annotations)
        gen_annotations = np.row_stack(gen_annotations) # ! 把所有label拼接起来
        sub_images.append(gen_image)
        sub_annotations.append(gen_annotations)

    #! 进行图片和标注替换
    for si,sub_image,sub_annotation in zip(sub_indexs,sub_images,sub_annotations):
        batch_images[si] = sub_image
        batch_annotations[si] = sub_annotation

    return batch_images, batch_annotations